import sys
import psycopg2
import db_conf

#CH2BrI would be 1*mass(C) + 2*mass(H) + 1*(Br) + 1*(I)
formula_file = open("iso.csv", "r")
isomass_hash = {}
for line in formula_file.readlines():
	tok = line.strip().split("\t")
	isomass_hash[tok[0]] = float(tok[1])
formula_file.close()

atomic_file = open("atomic_weights.csv", "r")
atomicweight_hash = {}
for line in atomic_file.readlines():
        tok = line.strip().split(",")
        atomicweight_hash[tok[0]] = float(tok[1])
atomic_file.close()

error_rpt = open('key_errors.log', 'w')
missing_elements = {}

def getElementMass(element, count):
	#print str(count) + "_" + element
	return isomass_hash[element]*count

def getAtomicWeight(element, count):
	try:
		#print str(count) + "_" + element
		return atomicweight_hash[element]*count	
	except KeyError,e:
		error_rpt.write(str(e) + "\n")
		if (missing_elements.get(str(e)) == None):
			missing_elements[element] = 1
		return 0;
 	
select_formula = "select compound_id, sum_formula from compound" #--monoisotopic_mass
config = db_conf.getConfig('metabvis.conf')
con = psycopg2.connect(database=db_conf.getDBName(config), port=db_conf.getDBPort(config), user=db_conf.getUser(config), password=db_conf.getPassword(config))
cur = con.cursor()
cur.execute(select_formula)
rows = cur.fetchall()
for r in rows:
	if (r[1] == None):
		continue
	print "id " + str(r[0]) + " formula " + r[1]
	#for each formula, break down the element, count and look up mass
	formula = r[1]
	id = r[0]
	actingElement = ""
	#for i in range(len(formula)):
	totalMass = 0.0
	totalWeight = 0.0
	i = 0	
	while i < len(formula):
		#print i
		if (i+1 < len(formula) and formula[i].isupper() == True and formula[i+1].islower() == False):
			if ((i + 1) < len(formula) and formula[i+1].isupper()):
				#rule 1, prior element has a count of 1
				totalMass = totalMass + getElementMass(formula[i], 1)
				totalWeight = totalWeight + getAtomicWeight(formula[i], 1)
				i = i + 1
			elif (formula[i+1].isdigit() and (i + 1) < len(formula)):
				if (i+3 < len(formula) and formula[i+2].isdigit() and formula[i+3].isdigit()):
					totalMass = totalMass + getElementMass(formula[i], int(formula[i+1:i+4]))
					totalWeight = totalWeight + getAtomicWeight(formula[i], int(formula[i+1:i+4]))
                                        i = i + 4
				elif (i+2 < len(formula) and formula[i+2].isdigit()):
					totalMass = totalMass + getElementMass(formula[i], int(formula[i+1:i+3]))
					totalWeight = totalWeight + getAtomicWeight(formula[i], int(formula[i+1:i+3]))
					i = i + 3
				else:
					totalMass = totalMass + getElementMass(formula[i], int(formula[i+1]))
					totalWeight = totalWeight + getAtomicWeight(formula[i], int(formula[i+1]))
					i = i + 2
			actingElement = ""
		else:
			#partial
			if (i+1 == len(formula)):
				totalMass = totalMass + getElementMass(formula[i], 1)
				totalWeight = totalWeight + getAtomicWeight(formula[i], 1)
				break
			if (i+2 == len(formula) and formula[-1].islower() == True):
				totalMass = totalMass + getElementMass(formula[i:i+2], 1)
				totalWeight = totalWeight + getAtomicWeight(formula[i:i+2], 1)
				break 
			actingElement = formula[i] + formula[i+1]
			if ((i + 2) < len(formula) and formula[i+2].isdigit() == False ):
				totalMass = totalMass + getElementMass(formula[i:i+2], 1)
				totalWeight = totalWeight + getAtomicWeight(formula[i:i+2], 1)
				i = i + 2				
			
			elif (formula[i+2].isdigit()):
				if (i + 3 < len(formula) and formula[i+3].isdigit()):
					totalMass = totalMass + getElementMass(formula[i:i+2], int(formula[i+2:i+4]))
					totalWeight = totalWeight + getAtomicWeight(formula[i:i+2], int(formula[i+2:i+4]))
					i = i + 4
				else:
					totalMass = totalMass + getElementMass(formula[i:i+2], int(formula[i+2])) 	
					totalWeight = totalWeight + getAtomicWeight(formula[i:i+2], int(formula[i+2]))
					i = i + 3 
	print formula + " monoisotopic_mass " + str(totalMass) + " molecular weight " + str(totalWeight)
	update = "update compound set monoisotopic_mass = %f, molecular_weight = %f where compound_id = %i" %(totalMass, totalWeight, id)
	con.cursor().execute(update)
	con.cursor().execute("commit")
print str(missing_elements)
#if __name__ == "__main__":
#        configfile = sys.argv[1]
#        print getDBHost(getConfig(configfile))

